package com.own.component.util.sensitive;

import com.own.component.util.sensitive.model.TrieNode;

import java.util.ArrayList;
import java.util.List;

/**
 * SensitiveWordFilter
 *
 * @author chenxueli
 * @date 2023-05-24 21:05:00
 */
public class SensitiveWordFilter {

    /**
     * 根节点
     */
    private final TrieNode root;
    /**
     * 替换敏感词的字符
     */
    private final char maskChar;

    public SensitiveWordFilter() {
        this('*');
    }

    public SensitiveWordFilter(char maskChar) {
        root = new TrieNode();
        this.maskChar = maskChar;
    }

    /**
     * 添加词语的方法
     *
     * @param word 词语
     */
    public void add(String word) {
        var node = root;
        for (var c : word.toCharArray()) {
            node.getChildren().putIfAbsent(c, new TrieNode());
            node = node.getChildren().get(c);
        }
        node.setEnd(true);
    }

    /**
     * 移除词语的方法
     *
     * @param word 词语
     */
    public void remove(String word) {
        var node = root;
        for (var ch : word.toCharArray()) {
            if (node.getChildren().containsKey(ch)) {
                node = node.getChildren().get(ch);
            } else {
                // 敏感词不存在，无需删除
                return;
            }
        }
        node.setEnd(false);
        if (node.getChildren().isEmpty()) {
            // 逐级向上删除没有子节点的父节点
            var parent = getParentNode(root, word);
            if (parent != null) {
                parent.getChildren().remove(word.charAt(word.length() - 1));
            }
        }
    }

    /**
     * 获取父级的节点
     *
     * @param node 当前节点
     * @param word 敏感词
     * @return 父级节点
     */
    private TrieNode getParentNode(TrieNode node, String word) {
        TrieNode parent = null;
        for (var i = 0; i < word.length() - 1; i++) {
            var ch = word.charAt(i);
            if (node.getChildren().containsKey(ch)) {
                parent = node;
                node = node.getChildren().get(ch);
            } else {
                return null;
            }
        }
        return parent;
    }

    /**
     * 是否包含敏感词
     *
     * @param text 文本
     * @return 是否包含敏感词
     */
    public boolean containsSensitiveWords(String text) {
        TrieNode node = root;
        for (var ch : text.toCharArray()) {
            if (node.getChildren().containsKey(ch)) {
                node = node.getChildren().get(ch);
                if (node.isEnd()) {
                    return true;
                }
            } else {
                node = root;
            }
        }
        return false;
    }

    /**
     * 使用指定的字符替换敏感词
     *
     * @param text 文本
     * @return 替换后的文本
     */
    public String replaceSensitiveWords(String text) {
        var filteredText = new StringBuilder(text);
        var sensitiveWords = getSensitiveWords(text);
        for (var word : sensitiveWords) {
            var startIndex = filteredText.indexOf(word);
            var endIndex = startIndex + word.length();
            filteredText.replace(startIndex, endIndex, String.valueOf(maskChar).repeat(word.length()));
        }
        return filteredText.toString();
    }

    /**
     * 获取敏感词
     *
     * @param text 文本
     * @return 敏感词列表
     */
    public List<String> getSensitiveWords(String text) {
        var sensitiveWords = new ArrayList<String>();
        var node = root;
        var start = -1;
        var end = -1;
        for (var i = 0; i < text.length(); i++) {
            var ch = text.charAt(i);
            if (node.getChildren().containsKey(ch)) {
                if (start == -1) {
                    start = i;
                }
                node = node.getChildren().get(ch);
                if (node.isEnd()) {
                    end = i + 1;
                    sensitiveWords.add(text.substring(start, end));
                    start = -1;
                    node = root;
                }
            } else {
                if (start != -1) {
                    start = -1;
                    node = root;
                }
            }
        }
        if (start != -1) {
            sensitiveWords.add(text.substring(start, end));
        }
        return sensitiveWords;
    }

}
